{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 7.5 PDE-Constrained Shape Optimization (fully automated)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want to solve the PDE-constrained shape optimization problem\n",
    "$$\n",
    "            \\underset{\\Omega\\subset \\mathsf{D}}{\\mbox{min}} \\; J(u) := \\int_\\Omega |u-u_d|^q \\; dx, \\quad q\\ge 2\n",
    "$$\n",
    "subject to that $(\\Omega,u)$ satisfy\n",
    "$$\n",
    "           \\int_\\Omega \\nabla u \\cdot \\nabla v \\; dx = \\int_\\Omega f v \\; dx \\; \\quad \\text{ for all } v \\in H_0^1(\\Omega),\n",
    "$$\n",
    "where $\\Omega \\subset \\mathbb R^2$ for given $u_d, f \\in C^1(\\mathbb R^2)$.\n",
    "\n",
    "Again, we want to compute the shape derivative by differentiation of a suitably defined perturbed Lagrangian using automated differentiation. Here this is accounted for automatically by NGSolve. For details we refer to \n",
    "\n",
    "P. Gangl, K. Sturm, M. Neunteufel, J. Schöberl.\n",
    "Fully and Semi-Automated Shape Differentiation in NGSolve,\n",
    "Struct. Multidisc. Optim., 63, pp.1579-1607, 2021."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve import *\n",
    "from ngsolve.webgui import Draw\n",
    "from netgen.geom2d import SplineGeometry\n",
    "from ngsolve.solvers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo = SplineGeometry()\n",
    "geo.AddCircle(c=(0.5,0.5), r=0.5, bc = 'circle')\n",
    "mesh = Mesh(geo.GenerateMesh(maxh = 0.08))\n",
    "Draw(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#given data of our problem (chosen such that \\Omega^* = [0,1]^2 is the optimal shape)\n",
    "f = CoefficientFunction(2*y*(1-y)+2*x*(1-x))\n",
    "ud = x*(1-x)*y*(1-y)\n",
    "\n",
    "grad_f = CoefficientFunction( (f.Diff(x), f.Diff(y) ) )\n",
    "grad_ud = CoefficientFunction( (ud.Diff(x), ud.Diff(y) ) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fes = H1(mesh, order=2, dirichlet=\".*\")\n",
    "gfu = GridFunction(fes)\n",
    "scene_u = Draw (gfu, mesh, \"state\")\n",
    "\n",
    "gfp = GridFunction(fes)\n",
    "scene_p = Draw (gfp, mesh, \"adjoint\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that (for linear problems) the operator on the left hand side of the adjoint equation is just the transpose of the state operator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Automatic Shape Differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The formula for the shape derivative was derived as the partial derivative of the perturbed Lagrangian (brought back to the original domain):\n",
    "$$\n",
    "    d\\mathcal J(\\Omega; X) = \\frac{\\partial}{\\partial t} \\left. \\left( \\int_\\Omega |u - u_d^t|^q  \\,\\mbox{det}(F_t) \\mbox{d} x \n",
    "         +  \\int_{\\Omega} (F_t^{-\\top}\\nabla u) \\cdot (F_t^{-\\top} \\nabla p)\\,  \\mbox{det}(F_t) \\, dx - \\int_{\\Omega}f^t p  \\,\\mbox{det}(F_t) \\,dx \\right) \\right \\rvert_{t=0} \n",
    "$$\n",
    "where \n",
    "<ul>\n",
    "    <li>   $T_t(x)=x+tX(x)=y$ \n",
    "    <li> $F_t = DT_t = I+t DX$\n",
    "    <li>  $u_d^t = u_d \\circ T_t$\n",
    "    <li> $f^t = f \\circ T_t$\n",
    "</ul>\n",
    "\n",
    "The integrand depends on the parameter $t$ only via $F_t$ and $T_t$. We define the Lagrangian in this form, involving a parameter t that has the value 0."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### In order to define the perturbed Lagrangian for a given problem, one needs to know transformation rules of differential operators, e.g., $$(\\nabla u)\\circ T_t = (\\partial T_t)^{-T} \\nabla (u \\circ T_t)$$ for $u \\in H^1$.\n",
    "### Since these transformation rules are known by NGSolve for all implemented differential operators, also the step of defining the perturbed Lagrangian can be automated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "VEC = H1(mesh, order=2, dim=2)\n",
    "PHI, X = VEC.TnT()\n",
    "\n",
    "def EquationFA(u,v):\n",
    "    return ( grad(u)*grad(v)-f*v)*dx\n",
    "\n",
    "q=4\n",
    "def CostAutoFA(u): \n",
    "    return (u-ud)**q*dx\n",
    "\n",
    "def CostAuto2(u): \n",
    "    return CostAutoFA(u)\n",
    "\n",
    "LagrangianFA = CostAutoFA(gfu) + EquationFA(gfu,gfp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### State equation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Equation can also be used to define the bilinear form. The following defines left and right hand side of the PDE in a \"BilinearForm\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u, v = fes.TnT()\n",
    "\n",
    "aAuto = BilinearForm(fes, symmetric=True)\n",
    "aAuto += EquationFA(u,v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the PDE can be conveniently solved by calling Newton's method (which terminates after one iteration since the PDE is linear)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfu.vec[:]=0\n",
    "Newton(aAuto,gfu,freedofs=fes.FreeDofs())\n",
    "scene_u.Redraw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adjoint equation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We set up the adjoint equation\n",
    "$$\n",
    "    \\mbox{Find } p \\in H_0^1(\\Omega): \\int_\\Omega \\nabla w \\cdot \\nabla p \\, \\mbox dx = - \\partial_u J(u)(w) \\quad \\text{ for all } w \\in H_0^1(\\Omega)\n",
    "$$\n",
    "where $u$ is the solution to the state equation. For $J(u) = \\int_\\Omega |u-u_d|^2 \\mbox dx$, we get\n",
    "$$\n",
    "    \\partial_u J(u)(w) = 2 \\int_\\Omega (u-u_d)w \\,\\mbox dx.\n",
    "$$\n",
    "However, we can also use the Diff(...) command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p, w = fes.TnT()\n",
    "\n",
    "fadjoint = LinearForm(fes)\n",
    "fadjoint += -1*(CostAutoFA(gfu)).Diff(gfu,w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def SolveAdjointEquation():\n",
    "    rhs = gfp.vec.CreateVector()\n",
    "    rhs.data = fadjoint.vec - aAuto.mat.T * gfp.vec\n",
    "    update = gfp.vec.CreateVector()\n",
    "    update.data = aAuto.mat.Inverse(fes.FreeDofs()).T * rhs\n",
    "    gfp.vec.data += update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fadjoint.Assemble()\n",
    "SolveAdjointEquation()\n",
    "scene_p.Redraw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Automatic Shape Differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Denoting the integrand by $G^{u,p}$, the shape derivative is given by\n",
    "$$\n",
    "\\begin{array}{rl}\n",
    "     d\\mathcal J(\\Omega; X) =& \\left( \\left. \\frac{\\partial G^{u,p}}{\\partial t} + \\frac{d  G^{u,p}}{dy} \\cdot \\frac{d T_t}{dt}\\right)\\right\\rvert_{t=0} \\\\\n",
    "     =& \\left.  \\frac{\\partial G^{u,p}}{\\partial t}\\right\\rvert_{t=0} + \\frac{d  G^{u,p}}{dy} \\cdot X\n",
    "\\end{array}\n",
    "$$\n",
    "The command .DiffShape(...) computes the shape derivative by automatically accounting for the corresponding transformations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dJOmegaAuto = LinearForm(VEC)\n",
    "dJOmegaAuto += LagrangianFA.DiffShape(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = BilinearForm(VEC)\n",
    "b += InnerProduct(grad(X),grad(PHI))*dx + InnerProduct(X,PHI)*dx\n",
    "\n",
    "gfX = GridFunction(VEC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def SolveDeformationEquationAuto():\n",
    "    rhs = gfX.vec.CreateVector()\n",
    "    rhs.data = dJOmegaAuto.vec - b.mat * gfX.vec\n",
    "    update = gfX.vec.CreateVector()\n",
    "    update.data = b.mat.Inverse(VEC.FreeDofs()) * rhs\n",
    "    gfX.vec.data += update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b.Assemble()\n",
    "dJOmegaAuto.Assemble()\n",
    "SolveDeformationEquationAuto()\n",
    "Draw(-gfX, mesh, \"-gfX\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gfset denotes the deformation of the original domain and will be updated during the shape optimization\n",
    "gfset = GridFunction(VEC)\n",
    "gfset.Set((0,0))\n",
    "mesh.SetDeformation(gfset)\n",
    "sceneSet = Draw(gfset,mesh,\"gfset\")\n",
    "SetVisualization (deformation=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfset.Set((0,0))\n",
    "mesh.SetDeformation(gfset)\n",
    "print('Cost at initial design', Integrate (CostAuto2(gfu), mesh))\n",
    "\n",
    "scale = 0.5 / Norm(gfX.vec)\n",
    "gfset.vec.data -= scale * gfX.vec\n",
    "sceneSet.Redraw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfu.vec[:]=0\n",
    "Newton(aAuto, gfu, fes.FreeDofs())\n",
    "print('Cost at new design', Integrate (CostAuto2(gfu), mesh))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus, the user has to enter the PDE (in its transformed form) only once."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let us again run the full algorithm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scene_u = Draw(gfu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#reset to and solve for initial configuration\n",
    "gfset.Set((0,0))\n",
    "mesh.SetDeformation(gfset)\n",
    "scene_u.Redraw()\n",
    "gfu.vec[:]=0\n",
    "Newton(aAuto, gfu, fes.FreeDofs())\n",
    "\n",
    "LineSearch = False\n",
    "\n",
    "\n",
    "iter_max = 600\n",
    "Jold = Integrate(CostAuto2(gfu), mesh)\n",
    "converged = False\n",
    "for k in range(iter_max):\n",
    "    print('cost at iteration', k, ': ', Jold)\n",
    "    mesh.SetDeformation(gfset)\n",
    "    scene_u.Redraw()\n",
    "    \n",
    "    gfu.vec[:]=0\n",
    "    Newton(aAuto, gfu, fes.FreeDofs(), printing = False)\n",
    "    \n",
    "    fadjoint.Assemble()\n",
    "    SolveAdjointEquation()\n",
    "    \n",
    "    b.Assemble()\n",
    "    dJOmegaAuto.Assemble()\n",
    "    SolveDeformationEquationAuto()\n",
    "\n",
    "    scale = 0.01 / Norm(gfX.vec)\n",
    "    gfsetOld = gfset\n",
    "    gfset.vec.data -= scale * gfX.vec\n",
    "    \n",
    "    Jnew = Integrate(CostAuto2(gfu), mesh)\n",
    "    \n",
    "    if LineSearch:\n",
    "        while Jnew > Jold and scale > 1e-12:\n",
    "            # input('a')\n",
    "            scale = scale / 2\n",
    "            \n",
    "            if scale <= 1e-12:\n",
    "                converged = True\n",
    "                break\n",
    "\n",
    "            gfset.vec.data = gfsetOld.vec - scale * gfX.vec\n",
    "            mesh.SetDeformation(gfset)\n",
    "            \n",
    "            gfu.vec[:]=0\n",
    "            Newton(aAuto, gfu, fes.FreeDofs(), printing = False)\n",
    "            Jnew = Integrate(CostAuto(gfu), mesh)\n",
    "    \n",
    "    if converged==True:\n",
    "        break\n",
    "    Jold = Jnew\n",
    "\n",
    "    Redraw(blocking=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
